# -*- coding: utf-8 -*-
# created on 2018/05/18

from mathsolver.functions.base import *
from sympy import Intersection, Union


class GongBiRangeUpdate001(BaseFunction):
    """
    已知等比数列{a_{n}}的公比为正数,且a_{5}•a_{7}=4{a_{4}}^{2},a_{2}=1,则a_{1}=( )
    """
    def solver(self, *args):
        assert len(args) == 2
        assert isinstance(args[0], BaseSequence)
        assert isinstance(args[1], str)
        sl = self.search(args[0].name)
        text = args[1]
        if text.find('正') >= 0:
            sl.qRange = Intersection(Interval(0, S.Infinity, True, True), sl.qRange)
        return self


class GongBiRangeUpdate002(BaseFunction):
    """
    eg. 已知公比不为1的等比数列{a_{n}}, a_{1}a_{2}a_{3}a_{4}a_{5}=243, 2a_{3}为3a_{2}和a_{4}的等差中项.求数列{a_{n}}的通项公式a_{n}.
    """
    def solver(self, *args):
        assert len(args) == 3
        assert isinstance(args[0], BaseSequence)
        sl = self.search(args[0].name)
        assert isinstance(args[1], str)
        text = args[1]
        assert isinstance(args[2], (BasePoly, BaseVariable, BaseNumber))
        value = args[2].sympify()
        if text.find("不") >= 0:
            q_range = Union(Interval(value, S.Infinity, True, True), Interval(- S.Infinity, value, True, True))
            sl.qRange = Intersection(sl.qRange, q_range)
        elif text.find("大于") >= 0:
            q_range = Interval(value, S.Infinity, True, True)
            sl.qRange = Intersection(sl.qRange, q_range)
        elif text.find("小于") >= 0:
            q_range = Interval(- S.Infinity, value, True, True)
            sl.qRange = Intersection(sl.qRange, q_range)
        self.output.append(sl)
        return self


class GongBiRangeUpdate003(BaseFunction):
    """
    在等比数列{a_{n}}中,a_{3}+a_{8}=-31,a_{4}a_{7}=-32,公比q是整数,求{a_{n}}的通项公式.
    """
    def solver(self, *args):
        assert len(args) == 2
        assert isinstance(args[0], BaseVariable)
        assert isinstance(args[1], str)
        sls = [self.known[item] for item in self.known if isinstance(self.known[item], BaseShuLieDengBi)]
        assert len(sls) == 1
        sl = sls[0]
        text = args[1]
        if text.find('整数') >= 0:
            sl.qValue = args[0].sympify()
            sl.qRange = range(1, 1000)
        self.output.append(sl)
        return self


class GongBiRangeUpdate004(BaseFunction):
    """
    设等比数列{a_{n}}的公比q<1,前n项和为S_{n}.已知a_{3}=2,S_{4}=5S_{2},求{a_{n}}的通项公式.
    """
    def solver(self, *args):
        assert len(args) == 2
        assert isinstance(args[0], BaseSequence)
        assert isinstance(args[1], BaseIneq)
        sl = self.search(args[0].name)
        ineq = args[1].sympify()
        expr = ineq[0] - ineq[2]
        symbols = list(expr.free_symbols)
        assert len(symbols) == 1
        q_value = symbols[0]
        sl.qValue = q_value
        self.known['inequations'].append(ineq)
        self.output.append(sl)
        return self


class GongBiRangeUpdate(BaseFunction):
    CLS = [GongBiRangeUpdate001, GongBiRangeUpdate002, GongBiRangeUpdate003, GongBiRangeUpdate004]

    def solver(self, *args):
        r = None
        for cl in GongBiRangeUpdate.CLS:
            try:
                r = cl(known=self.known, verbose=True).solver(*args)
                break
            except Exception:
                pass
        if not r:
            raise 'try fail'
        return r


class GongBiInitUpdate(BaseFunction):
    """
    等比数列公比不为零
    eg. 在等比数列{a_{n}}中,若a_{2}=4,a_{4}=1,则a_{6}=.
    """
    def solver(self, *args):
        assert len(args) == 1
        seq = args[0].sympify()
        answer = {}
        if 'Sequations' in self.known:
            seqs = self.search('Sequations')
        else:
            seqs = []
        p = r'[a-zA-Z]_\(.*?\)'
        m = re.findall(p, str(seq))
        m = set(m)
        m = list(m)
        assert len(m) == 1
        new_item = sympify(m[0])
        new_item_name = new_item.func
        sl = self.search(str(new_item_name))
        q_name = sl.qName
        self.steps.append(["", "%s" % BaseIneq([q_name, "!=", S.Zero]).printing()])
        seqs.append([q_name, "!=", S.Zero])
        answer['Sequations'] = seqs
        self.output.append(BaseShuLieValue(answer))
        return self
